# -*- coding: utf-8 -*-

import json
from bunch import Bunch


print ('Loading annotations and questions...')

ANNOTATIONS = 'annotations'
MULTIPLE_CHOICE_ANSWER = 'multiple_choice_answer'
QUESTION_ID = 'question_id'
QUESTION = 'question'
IMAGE_ID = 'image_id'
MULTIPLE_CHOICES = 'multiple_choices'
IMAGE_PATH = 'image_path'


def get_data_from_json(path):
    with open(path, 'r') as config_file:
        data_dict = json.load(config_file)

    # convert the dictionary to a namespace using bunch lib
    data = Bunch(data_dict)
    return data



def main():

    config_path = '../configs/prepro.json'
    config = get_data_from_json(config_path)

    root_path = config.root_path

    img_path = '%s/COCO_%s_%012d.jpg'

    train_ans_path = '%s/annotations/mscoco_train2014_annotations.json' % (root_path)
    val_ans_path = '%s/annotations/mscoco_val2014_annotations.json' % (root_path)

    train_ques_path = '%s/annotations/MultipleChoice_mscoco_train2014_questions.json' % (root_path)
    val_ques_path = '%s/annotations/MultipleChoice_mscoco_val2014_questions.json' % (root_path)


    train_ans= get_data_from_json(train_ans_path)
    train_qus = get_data_from_json(train_ques_path)

    val_ans = get_data_from_json(val_ans_path)
    val_qus = get_data_from_json(val_ques_path)

    # subtype = 'train2014'

    def load_ans(ans, ques, subtype):
        ans_list = []
        ans = ans.annotations
        ans_size = len(ans)
        for i in range(ans_size):
            ann = ans[i][MULTIPLE_CHOICE_ANSWER]
            ques_id = ans[i][QUESTION_ID]
            image_path = img_path % (subtype, subtype, ans[i][IMAGE_ID])

            question = ques.questions[i][QUESTION]
            mc_ans = ques.questions[i][MULTIPLE_CHOICES]

            ans_list.append({'ques_id': ques_id, 'img_path': image_path,
                             'question': question, 'MC_ans': mc_ans, 'ans': ann})
        return ans_list

    train = load_ans(train_ans, train_qus, subtype='train2014')
    test = load_ans(val_ans, val_qus, subtype='val2014')

    with open(config.vqa_raw_train, 'w') as f:
        json.dump(train, f)
    with open(config.vqa_raw_val, 'w') as f:
        json.dump(test, f)

    print('success!')

if __name__ == '__main__':
    main()
